import unittest
from main import extract_pets_from_part, extract_pets_from_cell, extract_pets_from_cell


class MyTestCase(unittest.TestCase):
    def test_extract_pets_from_part1(self):
        self.assertEqual({'阿拉斯加'}, extract_pets_from_part('萌宠_狗_阿拉斯加'))

    def test_extract_pets_from_part2(self):
        self.assertEqual({'阿拉斯加'}, extract_pets_from_part('阿拉斯加'))

    def test_extract_pets_from_cell1(self):
        self.assertEqual({'阿拉斯加'}, extract_pets_from_cell('萌宠_狗_阿拉斯加'))

    def test_extract_pets_from_cell1(self):
        self.assertEqual({'阿拉斯加', '边境牧羊犬'}, extract_pets_from_cell('萌宠_狗_阿拉斯加,萌宠_狗_边境牧羊犬'))


if __name__ == '__main__':
    unittest.main()
